#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
睡眠脑波数据处理系统 - 数据采集模块

文件功能：
- 封装NeuroPy3设备通信和数据采集
- 接收和处理原始EEG数据
- 监控信号质量和设备连接状态
- 提供数据回调接口给其他模块
- 集成日志记录功能

使用方法：
- 创建NeuroPyCollector实例，指定串口和日志管理器
- 设置数据回调函数处理接收到的数据
- 调用start_collection()开始数据采集
- 调用stop_collection()停止采集

原理：
- 基于NeuroPy3库与NeuroSky设备通信
- 使用回调机制实现数据的实时处理
- 监控信号质量，过滤低质量数据
- 将原始数据记录到日志系统中
"""

import sys
import time
import threading
from typing import Callable, Optional, Dict, Any
from pathlib import Path

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

try:
    from NeuroPy3 import NeuroPy
except ImportError:
    print("警告: 无法导入NeuroPy3，将使用模拟模式")
    NeuroPy = None

from .sleep_logger import SleepLogger


class NeuroPyCollector:
    """
    NeuroSky设备数据采集器
    
    负责与NeuroSky设备通信，采集原始EEG数据，并通过回调函数传递给其他模块
    """
    
    def __init__(self, port: str = "COM5", baudrate: int = 57600, logger: Optional[SleepLogger] = None):
        """
        初始化数据采集器
        
        Args:
            port: 串口号
            baudrate: 波特率
            logger: 日志管理器实例
        """
        self.port = port
        self.baudrate = baudrate
        self.logger = logger
        
        # 设备连接相关
        self.neuropy = None
        self.is_connected = False
        self.is_collecting = False
        
        # 数据回调函数
        self.data_callback = None
        self.quality_callback = None
        
        # 数据统计
        self.data_count = 0
        self.poor_signal_count = 0
        self.last_data_time = None
        
        # 线程控制
        self.collection_thread = None
        self.stop_event = threading.Event()
        
        # 模拟模式（当NeuroPy3不可用时）
        self.simulation_mode = NeuroPy is None
        
        if self.logger:
            self.logger.system_logger.info(f"初始化数据采集器: {port}@{baudrate}")
            if self.simulation_mode:
                self.logger.system_logger.warning("运行在模拟模式")
    
    def set_data_callback(self, callback: Callable[[Dict[str, Any]], None]):
        """
        设置数据回调函数
        
        Args:
            callback: 数据回调函数，接收包含EEG数据的字典
        """
        self.data_callback = callback
        if self.logger:
            self.logger.system_logger.info("设置数据回调函数")
    
    def set_quality_callback(self, callback: Callable[[int], None]):
        """
        设置信号质量回调函数
        
        Args:
            callback: 质量回调函数，接收信号质量值(0-200)
        """
        self.quality_callback = callback
        if self.logger:
            self.logger.system_logger.info("设置质量回调函数")
    
    def _on_raw_value(self, raw_value: int):
        """
        原始EEG数据回调处理
        
        Args:
            raw_value: 原始EEG数值
        """
        current_time = time.time()
        self.last_data_time = current_time
        self.data_count += 1
        
        # 获取当前的注意力和冥想值
        attention = getattr(self.neuropy, 'attention', 0) if self.neuropy else 0
        meditation = getattr(self.neuropy, 'meditation', 0) if self.neuropy else 0
        
        # 构建数据包
        data_packet = {
            'timestamp': current_time,
            'eeg_value': raw_value,
            'signal_quality': getattr(self.neuropy, 'poorSignal', 0) if self.neuropy else 0,
            'attention': attention,
            'meditation': meditation,
            'data_count': self.data_count
        }
        
        # 记录到日志
        if self.logger:
            self.logger.log_raw_data(data_packet)
        
        # 调用用户回调函数
        if self.data_callback:
            try:
                self.data_callback(data_packet)
            except Exception as e:
                if self.logger:
                    self.logger.system_logger.error(f"数据回调函数执行失败: {e}")
    
    def _on_poor_signal(self, poor_signal: int):
        """
        信号质量回调处理
        
        Args:
            poor_signal: 信号质量值 (0=良好, 200=无信号)
        """
        if poor_signal > 0:
            self.poor_signal_count += 1
        
        # 记录信号质量变化
        if self.logger:
            if poor_signal == 0:
                self.logger.system_logger.info("信号质量良好")
            elif poor_signal == 200:
                self.logger.system_logger.warning("信号丢失")
            else:
                self.logger.system_logger.warning(f"信号质量差: {poor_signal}")
        
        # 调用用户质量回调函数
        if self.quality_callback:
            try:
                self.quality_callback(poor_signal)
            except Exception as e:
                if self.logger:
                    self.logger.system_logger.error(f"质量回调函数执行失败: {e}")
    
    def _simulation_thread(self):
        """
        模拟数据生成线程（当真实设备不可用时）
        """
        import random
        
        if self.logger:
            self.logger.system_logger.info("开始模拟数据生成")
        
        data_count = 0
        while not self.stop_event.is_set():
            # 模拟原始EEG数据 (典型范围: -32768 到 32767)
            raw_value = random.randint(-20000, 20000)
            
            # 模拟信号质量 (大部分时间信号良好)
            poor_signal = random.choice([0, 0, 0, 0, 0, 200])  # 80%概率信号良好
            
            # 模拟注意力和冥想值
            attention = random.randint(20, 80)
            meditation = random.randint(15, 70)
            
            # 构建数据包
            data_packet = {
                'timestamp': time.time(),
                'eeg_value': raw_value,
                'signal_quality': poor_signal,
                'attention': attention,
                'meditation': meditation,
                'data_count': data_count
            }
            
            data_count += 1
            self.data_count = data_count
            
            # 记录到日志
            if self.logger:
                self.logger.log_raw_data(data_packet)
            
            # 调用回调函数
            if self.data_callback:
                try:
                    self.data_callback(data_packet)
                except Exception as e:
                    if self.logger:
                        self.logger.system_logger.error(f"模拟数据回调失败: {e}")
            
            # 模拟信号质量回调
            if self.quality_callback and random.random() < 0.1:  # 10%概率触发质量回调
                try:
                    self.quality_callback(poor_signal)
                except Exception as e:
                    if self.logger:
                        self.logger.system_logger.error(f"模拟质量回调失败: {e}")
            
            # 模拟512Hz采样率 (约2ms间隔)
            time.sleep(0.002)
        
        if self.logger:
            self.logger.system_logger.info("模拟数据生成结束")
    
    def connect(self) -> bool:
        """
        连接到NeuroSky设备
        
        Returns:
            bool: 连接是否成功
        """
        if self.simulation_mode:
            self.is_connected = True
            if self.logger:
                self.logger.system_logger.info("模拟设备连接成功")
            return True
        
        try:
            if self.logger:
                self.logger.system_logger.info(f"尝试连接设备: {self.port}")
            
            # 创建NeuroPy实例
            self.neuropy = NeuroPy(self.port, self.baudrate)
            
            # 设置回调函数
            self.neuropy.setCallBack("rawValue", self._on_raw_value)
            self.neuropy.setCallBack("poorSignal", self._on_poor_signal)
            
            # 开始连接
            self.neuropy.start()
            
            # 等待连接建立
            time.sleep(2)
            
            self.is_connected = True
            if self.logger:
                self.logger.system_logger.info("设备连接成功")
            
            return True
            
        except Exception as e:
            self.is_connected = False
            if self.logger:
                self.logger.system_logger.error(f"设备连接失败: {e}")
            return False
    
    def disconnect(self):
        """
        断开设备连接
        """
        if self.simulation_mode:
            self.is_connected = False
            if self.logger:
                self.logger.system_logger.info("模拟设备断开连接")
            return
        
        try:
            if self.neuropy:
                self.neuropy.stop()
                self.neuropy = None
            
            self.is_connected = False
            if self.logger:
                self.logger.system_logger.info("设备断开连接")
                
        except Exception as e:
            if self.logger:
                self.logger.system_logger.error(f"断开连接时发生错误: {e}")
    
    def start_collection(self) -> bool:
        """
        开始数据采集
        
        Returns:
            bool: 是否成功开始采集
        """
        if self.is_collecting:
            if self.logger:
                self.logger.system_logger.warning("数据采集已在进行中")
            return True
        
        if not self.is_connected:
            if not self.connect():
                return False
        
        # 重置统计信息
        self.data_count = 0
        self.poor_signal_count = 0
        self.last_data_time = time.time()
        self.stop_event.clear()
        
        # 启动采集
        self.is_collecting = True
        
        if self.simulation_mode:
            # 启动模拟线程
            self.collection_thread = threading.Thread(target=self._simulation_thread)
            self.collection_thread.daemon = True
            self.collection_thread.start()
        
        if self.logger:
            self.logger.system_logger.info("开始数据采集")
        
        return True
    
    def stop_collection(self):
        """
        停止数据采集
        """
        if not self.is_collecting:
            return
        
        self.is_collecting = False
        self.stop_event.set()
        
        # 等待模拟线程结束
        if self.collection_thread and self.collection_thread.is_alive():
            self.collection_thread.join(timeout=1.0)
        
        if self.logger:
            self.logger.system_logger.info(
                f"停止数据采集 - 总数据量: {self.data_count}, "
                f"信号质量差次数: {self.poor_signal_count}"
            )
    
    def get_signal_quality(self) -> Dict[str, Any]:
        """
        获取当前信号质量信息
        
        Returns:
            Dict[str, Any]: 信号质量统计信息
        """
        current_time = time.time()
        
        quality_info = {
            'is_connected': self.is_connected,
            'is_collecting': self.is_collecting,
            'data_count': self.data_count,
            'poor_signal_count': self.poor_signal_count,
            'signal_quality_ratio': 0.0,
            'last_data_time': self.last_data_time,
            'data_timeout': False
        }
        
        # 计算信号质量比例
        if self.data_count > 0:
            quality_info['signal_quality_ratio'] = 1.0 - (self.poor_signal_count / self.data_count)
        
        # 检查数据超时
        if self.last_data_time:
            time_since_last_data = current_time - self.last_data_time
            quality_info['data_timeout'] = time_since_last_data > 5.0  # 5秒无数据认为超时
        
        return quality_info
    
    def get_statistics(self) -> Dict[str, Any]:
        """
        获取采集统计信息
        
        Returns:
            Dict[str, Any]: 统计信息
        """
        quality_info = self.get_signal_quality()
        
        stats = {
            'port': self.port,
            'baudrate': self.baudrate,
            'simulation_mode': self.simulation_mode,
            'total_data_points': self.data_count,
            'poor_signal_events': self.poor_signal_count,
            'signal_quality_percentage': quality_info['signal_quality_ratio'] * 100,
            'connection_status': 'connected' if self.is_connected else 'disconnected',
            'collection_status': 'collecting' if self.is_collecting else 'stopped'
        }
        
        return stats
    
    def close(self):
        """
        关闭采集器，清理资源
        """
        if self.is_collecting:
            self.stop_collection()
        
        if self.is_connected:
            self.disconnect()
        
        if self.logger:
            stats = self.get_statistics()
            self.logger.system_logger.info(f"关闭数据采集器 - 统计信息: {stats}")


if __name__ == "__main__":
    # 简单测试
    print("测试NeuroPy数据采集器...")
    
    # 创建日志管理器
    logger = SleepLogger(log_dir="test_collector_logs", file_prefix="collector_test")
    session_id = logger.create_session_log("collector_test_001")
    
    # 创建采集器
    collector = NeuroPyCollector(port="COM5", logger=logger)
    
    # 设置数据回调
    def data_handler(data):
        print(f"接收数据: EEG={data['eeg_value']}, 质量={data['signal_quality']}, "
              f"注意力={data['attention']}, 冥想={data['meditation']}")
    
    def quality_handler(quality):
        if quality == 0:
            print("信号质量: 良好")
        elif quality == 200:
            print("信号质量: 无信号")
        else:
            print(f"信号质量: 差 ({quality})")
    
    collector.set_data_callback(data_handler)
    collector.set_quality_callback(quality_handler)
    
    try:
        # 开始采集
        if collector.start_collection():
            print("数据采集已开始，运行5秒...")
            time.sleep(5)
            
            # 显示统计信息
            stats = collector.get_statistics()
            print(f"\n统计信息: {stats}")
            
            # 停止采集
            collector.stop_collection()
            print("数据采集已停止")
        else:
            print("无法启动数据采集")
    
    except KeyboardInterrupt:
        print("\n用户中断")
    
    finally:
        # 清理资源
        collector.close()
        logger.close()
        print("测试完成")